#include <stdio.h>
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libavutil/avutil.h>
#include <libavutil/opt.h>
#include <libswresample/swresample.h>

int main(int argc, char *argv[]) {
    if (argc != 4) {
        printf("Usage: %s <input_file> <output_file> <quality(1-99)>\n", argv[0]);
        return 1;
    }

    const char *input_file = argv[1];
    const char *output_file = argv[2];
    int quality = atoi(argv[3]);

    if (quality < 1 || quality > 99) {
        printf("Quality must be between 1 and 99\n");
        return 1;
    }

    AVFormatContext *input_ctx = NULL;
    AVFormatContext *output_ctx = NULL;
    AVCodecContext *dec_ctx = NULL;
    AVCodecContext *enc_ctx = NULL;
    SwrContext *swr_ctx = NULL;
    AVFrame *frame = NULL;
    AVFrame *conv_frame = NULL;
    AVPacket *packet = NULL;
    AVPacket *enc_packet = NULL;
    int ret, stream_index = -1;

    // 打开输入文件
    if ((ret = avformat_open_input(&input_ctx, input_file, NULL, NULL)) < 0) {
        fprintf(stderr, "Could not open input file\n");
        goto end;
    }

    // 获取流信息
    if ((ret = avformat_find_stream_info(input_ctx, NULL)) < 0) {
        fprintf(stderr, "Could not find stream information\n");
        goto end;
    }

    // 找到音频流
    for (int i = 0; i < input_ctx->nb_streams; i++) {
        if (input_ctx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_AUDIO) {
            stream_index = i;
            break;
        }
    }

    if (stream_index == -1) {
        fprintf(stderr, "Could not find audio stream\n");
        goto end;
    }

    // 设置解码器
    const AVCodec *decoder = avcodec_find_decoder(input_ctx->streams[stream_index]->codecpar->codec_id);
    dec_ctx = avcodec_alloc_context3(decoder);
    if (!dec_ctx) {
        fprintf(stderr, "Could not allocate decoder context\n");
        goto end;
    }

    if ((ret = avcodec_parameters_to_context(dec_ctx, input_ctx->streams[stream_index]->codecpar)) < 0) {
        fprintf(stderr, "Could not copy decoder params\n");
        goto end;
    }

    if ((ret = avcodec_open2(dec_ctx, decoder, NULL)) < 0) {
        fprintf(stderr, "Could not open decoder\n");
        goto end;
    }

    // 创建输出上下文
    if ((ret = avformat_alloc_output_context2(&output_ctx, NULL, "mp3", output_file)) < 0) {
        fprintf(stderr, "Could not create output context\n");
        goto end;
    }

    // 设置编码器
    const AVCodec *encoder = avcodec_find_encoder(AV_CODEC_ID_MP3);
    AVStream *out_stream = avformat_new_stream(output_ctx, NULL);
    if (!out_stream) {
        fprintf(stderr, "Could not create output stream\n");
        goto end;
    }

    enc_ctx = avcodec_alloc_context3(encoder);
    if (!enc_ctx) {
        fprintf(stderr, "Could not allocate encoder context\n");
        goto end;
    }

    // 设置编码参数
    enc_ctx->sample_fmt = AV_SAMPLE_FMT_S16P;  // MP3编码器需要的采样格式
    enc_ctx->sample_rate = dec_ctx->sample_rate;
    enc_ctx->channel_layout = dec_ctx->channel_layout;
    enc_ctx->channels = dec_ctx->channels;
    // 根据质量参数调整比特率 (32-320kbps)
    enc_ctx->bit_rate = (32000 + (quality * 2880)); // 将1-99映射到32-320kbps
    enc_ctx->time_base = (AVRational){1, enc_ctx->sample_rate};

    // 设置一些MP3特定的参数
    if (encoder->capabilities & AV_CODEC_CAP_VARIABLE_FRAME_SIZE)
        enc_ctx->frame_size = 10000;

    // 将编码器参数复制到输出流
    if ((ret = avcodec_parameters_from_context(out_stream->codecpar, enc_ctx)) < 0) {
        fprintf(stderr, "Could not copy encoder params to output stream\n");
        goto end;
    }

    if ((ret = avcodec_open2(enc_ctx, encoder, NULL)) < 0) {
        fprintf(stderr, "Could not open encoder\n");
        goto end;
    }

    // 设置重采样上下文
    swr_ctx = swr_alloc_set_opts(NULL,
        enc_ctx->channel_layout,      // 输出布局
        enc_ctx->sample_fmt,          // 输出格式
        enc_ctx->sample_rate,         // 输出采样率
        dec_ctx->channel_layout,      // 输入布局
        dec_ctx->sample_fmt,          // 输入格式
        dec_ctx->sample_rate,         // 输入采样率
        0, NULL);

    if (!swr_ctx || swr_init(swr_ctx) < 0) {
        fprintf(stderr, "Could not initialize resampler\n");
        goto end;
    }

    // 打开输出文件
    if (!(output_ctx->oformat->flags & AVFMT_NOFILE)) {
        ret = avio_open(&output_ctx->pb, output_file, AVIO_FLAG_WRITE);
        if (ret < 0) {
            fprintf(stderr, "Could not open output file\n");
            goto end;
        }
    }

    // 写入文件头
    if ((ret = avformat_write_header(output_ctx, NULL)) < 0) {
        fprintf(stderr, "Could not write header\n");
        goto end;
    }

    // 分配帧和包
    frame = av_frame_alloc();
    conv_frame = av_frame_alloc();
    packet = av_packet_alloc();
    enc_packet = av_packet_alloc();

    if (!frame || !conv_frame || !packet || !enc_packet) {
        fprintf(stderr, "Could not allocate frame or packet\n");
        goto end;
    }

    // 读取、解码、编码循环
    while (av_read_frame(input_ctx, packet) >= 0) {
        if (packet->stream_index == stream_index) {
            if ((ret = avcodec_send_packet(dec_ctx, packet)) < 0) {
                fprintf(stderr, "Error sending packet to decoder\n");
                break;
            }

            while (ret >= 0) {
                ret = avcodec_receive_frame(dec_ctx, frame);
                if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) {
                    break;
                } else if (ret < 0) {
                    fprintf(stderr, "Error receiving frame from decoder\n");
                    goto end;
                }

                // 设置转换帧的参数
                conv_frame->sample_rate = enc_ctx->sample_rate;
                conv_frame->channel_layout = enc_ctx->channel_layout;
                conv_frame->format = enc_ctx->sample_fmt;
                conv_frame->nb_samples = frame->nb_samples;

                if ((ret = av_frame_get_buffer(conv_frame, 0)) < 0) {
                    fprintf(stderr, "Could not allocate output frame samples\n");
                    goto end;
                }

                // 重采样
                if ((ret = swr_convert(swr_ctx,
                    conv_frame->data, conv_frame->nb_samples,
                    (const uint8_t**)frame->data, frame->nb_samples)) < 0) {
                    fprintf(stderr, "Error while converting\n");
                    goto end;
                }

                conv_frame->pts = frame->pts;

                // 编码
                if ((ret = avcodec_send_frame(enc_ctx, conv_frame)) < 0) {
                    fprintf(stderr, "Error sending frame to encoder\n");
                    goto end;
                }

                while (ret >= 0) {
                    ret = avcodec_receive_packet(enc_ctx, enc_packet);
                    if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) {
                        break;
                    } else if (ret < 0) {
                        fprintf(stderr, "Error receiving packet from encoder\n");
                        goto end;
                    }

                    enc_packet->stream_index = 0;
                    av_packet_rescale_ts(enc_packet,
                        enc_ctx->time_base,
                        output_ctx->streams[0]->time_base);

                    ret = av_interleaved_write_frame(output_ctx, enc_packet);
                    if (ret < 0) {
                        fprintf(stderr, "Error writing packet\n");
                        goto end;
                    }
                    av_packet_unref(enc_packet);
                }
                av_frame_unref(frame);
                av_frame_unref(conv_frame);
            }
        }
        av_packet_unref(packet);
    }

    // 写入文件尾
    av_write_trailer(output_ctx);

end:
    // 清理资源
    if (dec_ctx) avcodec_free_context(&dec_ctx);
    if (enc_ctx) avcodec_free_context(&enc_ctx);
    if (frame) av_frame_free(&frame);
    if (conv_frame) av_frame_free(&conv_frame);
    if (packet) av_packet_free(&packet);
    if (enc_packet) av_packet_free(&enc_packet);
    if (swr_ctx) swr_free(&swr_ctx);
    if (input_ctx) avformat_close_input(&input_ctx);
    if (output_ctx) {
        if (!(output_ctx->oformat->flags & AVFMT_NOFILE))
            avio_closep(&output_ctx->pb);
        avformat_free_context(output_ctx);
    }

    return ret < 0 ? 1 : 0;
}
